# -*- coding: utf-8 -*-

import cv2
import math
import cc3d
import numpy as np


# 灰度线性拉伸
def stretch_linear(img_arr):
    return 255.0 * (img_arr - np.min(img_arr)) / (np.max(img_arr) - np.min(img_arr))


# 获取2D凸包络区域
def calc_convex_hull_2d(img_arr):
    # 1.获取输入区域所有像素坐标
    y_arr, x_arr = np.where(img_arr > 127)
    regions = np.concatenate((np.expand_dims(np.expand_dims(x_arr, axis=-1), axis=-1), 
                              np.expand_dims(np.expand_dims(y_arr, axis=-1), axis=-1)), axis=-1).astype(np.int32)
    
    # 2.计算凸包络
    hull_points = cv2.convexHull(regions)
    
    # 3.填充凸包络
    hull_arr = np.zeros_like(img_arr, dtype=np.uint8)
    cv2.fillPoly(hull_arr, [hull_points], 255)
    
    return hull_arr


# 根据均值mu_arr向量和协方差矩阵cov_arr构建2D高斯热力图
def put_heatmap_2d(heatmap, mu_arr, cov_arr):
    assert heatmap.ndim == 2
    assert mu_arr.shape == (2, 1) and cov_arr.shape == (2, 2)
    assert (cov_arr == cov_arr.T).all()
    th = 4.6052
    delta = math.sqrt(th * 2)
    # 使用 np.linalg.pinvh 直接计算协方差矩阵的伪逆
    cov_det = np.linalg.det(cov_arr)

    if cov_det == 0:
        cov_arr += 1e-12 * np.eye(cov_arr.shape[0], cov_arr.shape[1])
        cov_det = np.linalg.det(cov_arr)

    coinv_arr = np.linalg.inv(cov_arr)

    start_coord = np.maximum(mu_arr.squeeze(-1) - delta * np.sqrt(np.max(cov_arr, axis=-1)), 0).astype(np.int32)
    end_coord = np.minimum(mu_arr.squeeze(-1) + delta * np.sqrt(np.max(cov_arr, axis=-1)),
                           np.array(heatmap.shape[::-1])).astype(np.int32)

    xx, yy = np.meshgrid(np.arange(start_coord[1], end_coord[1]), np.arange(start_coord[0], end_coord[0]))
    pp = np.vstack([yy.ravel(), xx.ravel()])

    diffs = pp - mu_arr
    exponents = np.sum(np.dot(coinv_arr, diffs) * diffs, axis=0) / 2.0

    # 使用向量化运算计算有效点，避免后续的不必要计算
    valid_indices = exponents <= th
    valid_pp = pp[:, valid_indices]
    valid_exponents = exponents[valid_indices]
    # 使用向量化运算计算有效点的指数部分
    valid_values = np.exp(-valid_exponents)
    valid_values = np.maximum(np.minimum(valid_values, 1.0), 0.0)
    # 使用向量化运算更新热力图
    # 使用 valid_pp[1], valid_pp[0] 作为索引
    heatmap[valid_pp[1], valid_pp[0]] = np.maximum(heatmap[valid_pp[1], valid_pp[0]], valid_values)

    return heatmap


# 根据均值mu_arr向量和协方差矩阵cov_arr构建3D高斯热力图
def put_heatmap_3d(heatmap, mu_arr, cov_arr):
    assert heatmap.ndim == 3
    assert mu_arr.shape == (3, 1) and cov_arr.shape == (3, 3)
    assert (cov_arr == cov_arr.T).all()
    
    th = 4.6052
    delta = math.sqrt(th * 2)  # 3.0348640826238
    if np.linalg.det(cov_arr) == 0:  # 确保协方差矩阵cov_arr正定
        cov_arr += 1e-12 * np.eye(cov_arr.shape[0], cov_arr.shape[1])
    coinv_arr = np.linalg.inv(cov_arr)  # 求协方差矩阵cov_arr的逆矩阵
    
    # 限定计算热力图分布坐标范围
    start_coord = mu_arr.squeeze(-1) - delta * np.sqrt(np.max(cov_arr, axis=-1))
    start_coord[start_coord < 0] = 0
    start_coord = start_coord.astype(np.int32)
    
    end_coord = mu_arr.squeeze(-1) + delta * np.sqrt(np.max(cov_arr, axis=-1))
    end_coord = np.minimum(end_coord, np.array(heatmap.shape[::-1]))
    end_coord = end_coord.astype(np.int32)
    
    pp = np.mgrid[start_coord[0]:end_coord[0]:1, start_coord[1]:end_coord[1]:1, start_coord[2]:end_coord[2]:1]
    pp = np.reshape(pp, (3, -1))
    for idx in range(pp.shape[-1]):  # 遍历限定计算热力图分布范围的点计算热力值
        exp = np.dot(np.dot((pp[:, idx:idx+1] - mu_arr).T, coinv_arr), pp[:, idx:idx+1] - mu_arr).squeeze() / 2.0
        if exp > th:
            continue
        heatmap[pp[2, idx], pp[1, idx], pp[0, idx]] = max(heatmap[pp[2, idx], pp[1, idx], pp[0, idx]], math.exp(-exp))
        heatmap[pp[2, idx], pp[1, idx], pp[0, idx]] = min(heatmap[pp[2, idx], pp[1, idx], pp[0, idx]], 1.0)
    
    return heatmap


# 获取3D连通域信息
def calc_cc3d(img_arr):
    labels_out, N = cc3d.connected_components(img_arr, connectivity=26, return_N=True)
    ccinfo_dct = {"S": [], "mean": [], "cov": []}
    for idx in range(1, N+1):
        coord_arr = np.array(np.where(labels_out == idx), dtype=np.int32)[::-1, :].T
        ccinfo_dct["S"].append((labels_out == idx).sum())
        ccinfo_dct["mean"].append(np.expand_dims(np.mean(coord_arr, axis=0), axis=-1))
        ccinfo_dct["cov"].append(np.cov(coord_arr, rowvar=False))
    ccinfo_dct["S"].append(np.mean(np.array(ccinfo_dct["S"]), axis=0))
    ccinfo_dct["cov"].append(np.cov(np.array(ccinfo_dct["mean"]).squeeze(-1), rowvar=False))
    ccinfo_dct["mean"].append(np.mean(np.array(ccinfo_dct["mean"]), axis=0))
    
    return ccinfo_dct
